/*
 * Copyright 2025 The Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 *@file    : feature_matcher.h
 *@brief   : feature match 
 *@authors : zhanglei, zhanglei_723@126.com
 *@version : v1.0
 *@data    : 2025/8/30
 *
 */

#ifndef FEATURE_MATCHER_H
#define FEATURE_MATCHER_H

#include <Eigen/Eigen>
#include <iostream>
#include <unordered_map>
#include <vector>

#include "feature/feature_database.h"

namespace feature {

class FeatureMatcher {
public:
    FeatureMatcher(std::shared_ptr<FeatureDatabase> feature_database_ptr);
    void DescriptorMatch(uint64_t& feature_id, const uint8_t& camera_id, const uint64_t& frame_id, const uint64_t& last_frame_id, const std::vector<FeaturePoint>& feature_data);
    void OrbDescriptorMatch(uint64_t& feature_id, const uint8_t& camera_id, const uint64_t& frame_id, const uint64_t& last_frame_id, const std::vector<FeaturePoint>& feature_data);
    void RobustDescriptorMatch(uint64_t& feature_id, const uint8_t& camera_id, const uint64_t& frame_id, const uint64_t& last_frame_id, const std::vector<FeaturePoint>& feature_data);
private:
    inline double ComputeDescriptorDistance(const Vector128d &d1, const Vector128d &d2) {
        double distance = d1.dot(d2);
        distance = std::sqrt(2 - 2 * distance);
        return distance;
    }
    inline unsigned int ComputeDescriptorDistance(const cv::Mat& desc_1, const cv::Mat& desc_2) {
        // http://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetParallel

        constexpr uint32_t mask_1 = 0x55555555U;
        constexpr uint32_t mask_2 = 0x33333333U;
        constexpr uint32_t mask_3 = 0x0F0F0F0FU;
        constexpr uint32_t mask_4 = 0x01010101U;

        const auto* pa = desc_1.ptr<uint32_t>();
        const auto* pb = desc_2.ptr<uint32_t>();

        unsigned int dist = 0;

        for (unsigned int i = 0; i < 8; ++i, ++pa, ++pb) {
            auto v = *pa ^ *pb;
            v -= ((v >> 1) & mask_1);
            v = (v & mask_2) + ((v >> 2) & mask_2);
            dist += (((v + (v >> 4)) & mask_3) * mask_4) >> 24;
        }

        return dist;
    }

    
    void FilterByFundamentalMatrix(std::vector<uchar>& inliers, const std::vector<Vector3d>& last_point, const std::vector<Vector3d>& curr_point);

    /**
    * @brief Find matches between two keypoint+descriptor sets.
    * @param pts0 first vector of keypoints
    * @param pts1 second vector of keypoints
    * @param desc0 first vector of descriptors
    * @param desc1 second vector of decriptors
    * @param id0 id of the first camera
    * @param id1 id of the second camera
    * @param matches vector of matches that we have found
    *
    * This will perform a "robust match" between the two sets of points (slow but has great results).
    * First we do a simple KNN match from 1to2 and 2to1, which is followed by a ratio check and symmetry check.
    * Original code is from the "RobustMatcher" in the opencv examples, and seems to give very good results in the matches.
    * https://github.com/opencv/opencv/blob/master/samples/cpp/tutorial_code/calib3d/real_time_pose_estimation/src/RobustMatcher.cpp
    */
    void RobustMatch(const std::vector<cv::KeyPoint> &pts0, const std::vector<cv::KeyPoint> &pts1, const cv::Mat &desc0,
                    const cv::Mat &desc1, const std::vector<Vector2d>& last_undistorted_point, const std::vector<Vector2d> curr_undistorted_point, 
                    std::vector<cv::DMatch> &matches);

    // Helper functions for the robust_match function
    // Original code is from the "RobustMatcher" in the opencv examples
    // https://github.com/opencv/opencv/blob/master/samples/cpp/tutorial_code/calib3d/real_time_pose_estimation/src/RobustMatcher.cpp
    void RobustRatioTest(std::vector<std::vector<cv::DMatch>> &matches);
    void RobustSymmetryTest(std::vector<std::vector<cv::DMatch>> &matches1, std::vector<std::vector<cv::DMatch>> &matches2,
                            std::vector<cv::DMatch> &good_matches);

    // Our descriptor matcher
    cv::Ptr<cv::DescriptorMatcher> matcher_ = cv::DescriptorMatcher::create("BruteForce-Hamming");

    std::shared_ptr<FeatureDatabase> feature_database_ptr_;

    double window_size_; 
};

} // namespace feature

#endif /* FEATURE_MATCHER_H */